"""
协同过滤
"""

import os
import numpy as np
import pandas as pd


DATA_PATH = './ml-latest-small/ratings.csv'
CACHE_DIR = './cache'


def load_data(filepath):
    cache_path = os.path.join(CACHE_DIR, 'ratings_matrix.cache')
    print('开始加载数据集...')
    if os.path.exists(cache_path):  # 判断是否存在缓存文件
        print('缓存加载中...')
        ratings_matrix = pd.read_pickle(cache_path)
        print('从缓存加载数据集完成')
    else:
        print('新数据加载中...')
        # 设置要加载字段的数据类型
        dtype = {'userId': np.int32, 'movieId': np.int32, 'rating': np.float32}
        # 加载数据前三列
        ratings = pd.read_csv(filepath, dtype=dtype, usecols=range(3))
        # 透视表
        ratings_matrix = ratings.pivot_table(index=['userId'], columns=['movieId'], values='rating')
        # 保存缓存文件
        ratings_matrix.to_pickle(cache_path)
        print('从新数据加载数据集完成')
    return ratings_matrix


def compute_similarity(ratings_matrix, based='user'):
    """
    计算相似度 皮尔逊相关系数
    :param ratings_matrix: 用户-物品评分矩阵
    :param based: 可选值范围['user', 'item']
    :return:
    """

    user_similarity_cache_path = os.path.join(CACHE_DIR, 'user_similarity.cache')
    item_similarity_cache_path = os.path.join(CACHE_DIR, 'item_similarity.cache')

    # 用户相似度
    if based == 'user':
        if os.path.exists(user_similarity_cache_path):
            print('从缓存加载用户相似度...')
            similarity = pd.read_pickle(user_similarity_cache_path)
        else:
            print('开始计算用户相似度...')
            similarity = ratings_matrix.T.corr()
            similarity.to_pickle(user_similarity_cache_path)
    # 物品相似度
    elif based == 'item':
        if os.path.exists(item_similarity_cache_path):
            print('从缓存加载物品相似度...')
            similarity = pd.read_pickle(item_similarity_cache_path)
        else:
            print('开始计算物品相似度...')
            similarity = ratings_matrix.corr()
            similarity.to_pickle(item_similarity_cache_path)
    else:
        raise Exception('未知的based参数值：%s' % based)
    print('相似度计算/加载完成')
    return similarity


def user_based_predict(user_id, item_id, ratings_matrix, user_similarity):
    """
    预测指定用户对指定物品的评分值
    :param user_id:
    :param item_id:
    :param ratings_matrix: 用户-物品评分矩阵
    :param user_similarity: 用户相似度矩阵
    :return:
    """
    print('开始预测用户<%d>对电影<%d>的评分...' % (user_id, item_id))
    # 1. 找出用户的相似用户
    similar_users = user_similarity[user_id].drop([user_id]).dropna()
    # 相似用户筛选规则：正相关用户
    similar_users = similar_users.where(similar_users > 0).dropna()
    if similar_users.empty is True:
        raise Exception('用户<%d>没有相似的用户' % user_id)
    # 2. 从用户的近邻相似用户中筛选出对物品有评分记录的用户
    ids = set(ratings_matrix[item_id].dropna().index) & set(similar_users.index)  # 交集
    finally_similar_users = similar_users.loc[list(ids)]
    # 3. 结合用户与其近邻用户的相似度预测用户对物品的评分
    sum_up = 0  # 评分公式的分子部分
    sum_down = 0  # 评分公式的分母部分
    for sim_user_id, similarity in finally_similar_users.iteritems():
        # 近邻用户的评分数据
        sim_user_rated_movies = ratings_matrix.loc[sim_user_id].dropna()
        # 近邻用户对物品的评分
        sim_user_rating_for_item = sim_user_rated_movies[item_id]
        sum_up += similarity * sim_user_rating_for_item
        sum_down += similarity
    # 计算预测的评分值
    predicted_rating = sum_up / sum_down
    print('预测用户<%d>对电影<%d>的评分为：%.2f' % (user_id, item_id, predicted_rating))
    return round(predicted_rating, 2)


def user_based_predict_all(user_id, ratings_matrix, user_similarity):
    """
    预测全部评分
    :param user_id:
    :param ratings_matrix: 用户-物品评分矩阵
    :param user_similarity: 用户相似度矩阵
    :return:
    """
    # 要预测的物品列表
    item_ids = ratings_matrix.columns
    for item_id in item_ids:
        try:
            rating = user_based_predict(user_id, item_id, ratings_matrix, user_similarity)
        except Exception as e:
            print(e)
        else:
            yield user_id, item_id, rating


def user_based_top_k_result(user_id, k):
    """
    根据预测评分为指定用户进行TOP-N推荐
    :param user_id:
    :param k:
    :return:
    """
    ratings_matrix = load_data(DATA_PATH)
    user_similarity = compute_similarity(ratings_matrix)
    results = user_based_predict_all(user_id, ratings_matrix, user_similarity)
    return sorted(results, key=lambda x: x[2], reverse=True)[:k]


def item_based_predict(user_id, item_id, ratings_matrix, item_similarity):
    """
    预测指定用户对指定物品的评分值
    :param user_id:
    :param item_id:
    :param ratings_matrix: 用户-物品评分矩阵
    :param item_similarity: 物品相似度矩阵
    :return:
    """
    print('开始预测用户<%d>对电影<%d>的评分...' % (user_id, item_id))
    # 1. 找出物品的相似物品
    similar_items = item_similarity[item_id].drop([item_id]).dropna()
    # 相似物品筛选规则：正相关物品
    similar_items = similar_items.where(similar_items > 0).dropna()
    if similar_items.empty is True:
        raise Exception('物品<%d>没有相似的物品' % item_id)
    # 2. 从物品的近邻相似物品中筛选出用户评分过的物品
    ids = set(ratings_matrix.loc[user_id].dropna().index) & set(similar_items.index)  # 交集
    finally_similar_items = similar_items.loc[list(ids)]
    # 3. 结合物品与其近邻物品的相似度和用户对其相似物品的评分预测用户对物品的评分
    sum_up = 0  # 评分公式的分子部分
    sum_down = 0  # 评分公式的分母部分
    for sim_item_id, similarity in finally_similar_items.iteritems():
        # 近邻物品的评分数据
        sim_item_rated_movies = ratings_matrix[sim_item_id].dropna()
        # 用户对相似物品的评分
        sim_item_rating_from_user = sim_item_rated_movies[user_id]
        sum_up += similarity * sim_item_rating_from_user
        sum_down += similarity
    # 计算预测的评分值
    predicted_rating = sum_up / sum_down
    print('预测用户<%d>对电影<%d>的评分为：%.2f' % (user_id, item_id, predicted_rating))
    return round(predicted_rating, 2)


def item_based_predict_all(user_id, ratings_matrix, item_similarity):
    """
    预测全部评分
    :param user_id:
    :param ratings_matrix: 用户-物品评分矩阵
    :param item_similarity: 物品相似度矩阵
    :return:
    """
    # 要预测的物品列表
    item_ids = ratings_matrix.columns
    for item_id in item_ids:
        try:
            rating = item_based_predict(user_id, item_id, ratings_matrix, item_similarity)
        except Exception as e:
            print(e)
        else:
            yield user_id, item_id, rating


def item_based_top_k_result(user_id, k):
    """
    根据预测评分为指定用户进行TOP-N推荐
    :param user_id:
    :param k:
    :return:
    """
    ratings_matrix = load_data(DATA_PATH)
    item_similarity = compute_similarity(ratings_matrix, based='item')
    results = item_based_predict_all(user_id, ratings_matrix, item_similarity)
    return sorted(results, key=lambda x: x[2], reverse=True)[:k]


if __name__ == '__main__':
    import time
    start = time.time()
    print('=' * 50)
    results1 = user_based_top_k_result(1, 20)
    print(time.time()-start)
    start = time.time()
    print('=' * 50)
    results2 = item_based_top_k_result(1, 20)
    print(time.time() - start)
    from pprint import pprint
    pprint(results1)
    pprint(results2)
